package ws

import (
	"sync"

	"gitee.com/AlexWoo/arpc/alog"
	"gitee.com/AlexWoo/arpc/anet"
	"github.com/fasthttp/router"
	"github.com/fasthttp/websocket"
	"github.com/valyala/fasthttp"
)

type CreateHandler func(*Connection) anet.Handler

type Server struct {
	fs *fasthttp.Server

	conf *ServerConf

	lctx *alog.LogCtx // Log Context

	path   string
	create CreateHandler

	wg *sync.WaitGroup
}

func NewServer(c *ServerConf) *Server {
	logger := alog.NewWsLog(&c.Log)

	s := &Server{
		fs: &fasthttp.Server{
			Logger: logger,
		},
		conf: c,
		lctx: &alog.LogCtx{
			LogLevel: alog.Level[c.Log.LogLevel],
		},
	}

	return s
}

func (s *Server) RegisteHandler(path string, ch CreateHandler) {
	s.path = path
	s.create = ch
}

func (s *Server) Run(wg *sync.WaitGroup) {
	s.wg = wg

	r := router.New()
	r.GET(s.path, s.handle)
	s.fs.Handler = r.Handler

	wg.Add(1)
	go func(wg *sync.WaitGroup) {
		defer wg.Done()

		err := s.fs.ListenAndServe(s.conf.Addr)
		if err != nil {
			alog.LogError(s.lctx, "Error in ListenAndServe: %v", err)
		}

		alog.LogInfo(s.lctx, "Websocket Server Stop!!!!")
	}(wg)
}

func (s *Server) Stop() {
	s.fs.Shutdown()
}

func (s *Server) handle(ctx *fasthttp.RequestCtx) {
	upgrader := websocket.FastHTTPUpgrader{
		CheckOrigin: func(ctx *fasthttp.RequestCtx) bool { // 鉴权
			return true
		},
	}

	s.wg.Add(1)
	err := upgrader.Upgrade(ctx, s.createConnction)
	if err != nil {
		alog.LogError(s.lctx, "Handler websocket upgrade %v", err)
		s.wg.Done()
		return
	}
}

func (s *Server) createConnction(wc *websocket.Conn) {
	defer s.wg.Done()

	c := newConnection(&s.conf.Log)
	c.conn = wc
	c.h = s.create(c)
	c.remote = wc.RemoteAddr().String()
	c.local = wc.LocalAddr().String()

	alog.LogInfo(c.lctx, "WSServer accept new connection")

	c.loop()
}
